#!/usr/bin/env python3
"""
重新提取音频脚本
根据修改后的音频提取逻辑（前后各延长0.5秒），重新从原视频中提取音频片段并上传到服务器
"""

import os
import json
import sys
from pathlib import Path
from moviepy.editor import VideoFileClip
import uuid
import tempfile
import logging
import paramiko
from scp import SCPClient

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.append(str(project_root / "code"))

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class AudioReExtractor:
    def __init__(self):
        """初始化音频重新提取器"""
        # 服务器连接配置
        self.server_host = "159.138.23.13"
        self.server_username = "root"
        self.server_password = ""  # 请在这里填入服务器密码
        self.server_base_path = "/data"

        if not self.server_password:
            logging.error("请在代码中设置服务器密码")
            sys.exit(1)
        
    def parse_timestamp(self, timestamp_str):
        """
        解析时间戳字符串为开始和结束时间（毫秒）
        格式: "01:22:14,606 --> 01:22:16,846"
        """
        try:
            start_end = timestamp_str.split(' --> ')
            if len(start_end) != 2:
                raise ValueError(f"无效的时间戳格式: {timestamp_str}")

            start_time = start_end[0].strip()
            end_time = start_end[1].strip()

            def time_str_to_ms(time_str):
                hh_mm_ss, ms = time_str.split(',')
                h, m, s = hh_mm_ss.split(':')
                total_ms = int(h) * 3600 * 1000 + int(m) * 60 * 1000 + int(s) * 1000 + int(ms)
                return total_ms

            return time_str_to_ms(start_time), time_str_to_ms(end_time)
        except Exception as e:
            logging.error(f"解析时间戳失败: {timestamp_str}, 错误: {e}")
            return None, None

    def upload_to_server(self, local_file_path, remote_file_path):
        """
        上传文件到服务器
        """
        try:
            # 创建SSH连接
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(
                hostname=self.server_host,
                username=self.server_username,
                password=self.server_password,
                timeout=30
            )

            # 确保远程目录存在
            remote_dir = os.path.dirname(remote_file_path)
            stdin, stdout, stderr = ssh.exec_command(f"mkdir -p {remote_dir}")
            stdout.channel.recv_exit_status()  # 等待命令执行完成

            # 创建SCP客户端并上传文件
            with SCPClient(ssh.get_transport()) as scp:
                scp.put(local_file_path, remote_file_path)
                logging.info(f"成功上传文件到服务器: {remote_file_path}")

            ssh.close()
            return True

        except Exception as e:
            logging.error(f"上传文件到服务器失败: {e}")
            return False

    def extract_audio_segment(self, video_path, start_ms, end_ms, output_path, video_duration):
        """
        从视频中提取音频片段（应用新的延长逻辑）
        """
        try:
            # 应用新的音频提取逻辑：前后各延长0.5秒
            start_sec = start_ms / 1000  # 转换为秒
            start_sec = max(0, start_sec - 0.5)  # 向前延长0.5秒，但不能小于0
            end_sec = end_ms / 1000  # 转换为秒
            end_sec = min(end_sec + 0.5, video_duration)  # 向后延长0.5秒，但不能超过视频时长

            # 确保时间范围有效
            if start_sec >= end_sec:
                logging.error(f"无效的时间范围: {start_sec}s - {end_sec}s")
                return False

            # 加载视频并提取音频片段
            video = VideoFileClip(video_path)
            audio_clip = video.subclip(start_sec, end_sec).audio
            
            # 保存音频片段
            audio_clip.write_audiofile(output_path, codec='mp3', verbose=False, logger=None)
            
            # 清理资源
            audio_clip.close()
            video.close()
            
            logging.info(f"成功提取音频片段: {output_path} ({start_sec:.1f}s - {end_sec:.1f}s)")
            return True
            
        except Exception as e:
            logging.error(f"提取音频片段失败: {e}")
            return False

    def process_dataset_file(self, json_file_path, video_path):
        """
        处理单个数据集文件

        参数:
            json_file_path: JSON数据集文件路径
            video_path: 对应的视频文件路径
        """
        logging.info(f"开始处理数据集文件: {json_file_path}")
        logging.info(f"使用视频文件: {video_path}")

        try:
            # 验证视频文件是否存在
            if not Path(video_path).exists():
                logging.error(f"视频文件不存在: {video_path}")
                return False

            # 读取JSON文件
            with open(json_file_path, 'r', encoding='utf-8') as f:
                dataset = json.load(f)

            # 从文件名提取IP ID
            filename = Path(json_file_path).name
            ip_id = None
            # 尝试从文件名中提取IP ID
            if 'word_dataset_' in filename:
                parts = filename.split('_')
                for part in parts:
                    if part.startswith('IP'):
                        ip_id = part
                        break

            if not ip_id:
                logging.error(f"无法从文件名中提取IP ID: {filename}")
                return False
            
            # 获取视频时长
            video = VideoFileClip(video_path)
            video_duration = video.duration
            video.close()
            
            logging.info(f"视频文件: {video_path}, 时长: {video_duration:.1f}秒")
            
            # 统计信息
            total_items = 0
            processed_items = 0
            failed_items = 0
            
            # 遍历数据集中的所有单词
            for word, word_data_list in dataset.items():
                for word_data in word_data_list:
                    total_items += 1
                    
                    # 获取必要信息
                    timestamp = word_data.get('timestamp')
                    audio_id = word_data.get('audio_id')
                    audio_path = word_data.get('audio_path')
                    
                    if not all([timestamp, audio_id, audio_path]):
                        logging.warning(f"数据不完整，跳过: {word_data}")
                        failed_items += 1
                        continue
                    
                    # 解析时间戳
                    start_ms, end_ms = self.parse_timestamp(timestamp)
                    if start_ms is None or end_ms is None:
                        logging.warning(f"时间戳解析失败，跳过: {timestamp}")
                        failed_items += 1
                        continue
                    
                    # 创建临时音频文件
                    with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as temp_file:
                        temp_audio_path = temp_file.name
                    
                    try:
                        # 提取音频片段
                        if self.extract_audio_segment(video_path, start_ms, end_ms, temp_audio_path, video_duration):
                            # 构建服务器路径（保持原有路径结构）
                            # 从URL中提取路径: http://159.138.23.13/audio/original/IP1fef9dba/43c7dc48.mp3
                            # 提取: audio/original/IP1fef9dba/43c7dc48.mp3
                            if 'audio/original/' in audio_path:
                                relative_path = audio_path.split('audio/original/')[-1]
                                server_path = f"{self.server_base_path}/audio/original/{relative_path}"
                            else:
                                # 备用方案
                                server_path = f"{self.server_base_path}/audio/original/{ip_id}/{audio_id}.mp3"

                            # 上传到服务器
                            if self.upload_to_server(temp_audio_path, server_path):
                                logging.info(f"成功重新提取并上传音频: {audio_id} -> {server_path}")
                                processed_items += 1
                            else:
                                logging.error(f"上传音频失败: {audio_id}")
                                failed_items += 1
                        else:
                            logging.error(f"提取音频失败: {audio_id}")
                            failed_items += 1
                            
                    finally:
                        # 清理临时文件
                        if os.path.exists(temp_audio_path):
                            os.unlink(temp_audio_path)
            
            logging.info(f"处理完成 - 总计: {total_items}, 成功: {processed_items}, 失败: {failed_items}")
            return True
            
        except Exception as e:
            logging.error(f"处理数据集文件失败: {e}")
            return False



    def run(self, ip_dir, video_path):
        """
        运行重新提取音频的主流程

        参数:
            ip_dir: 要处理的IP目录名（如 "IP1fef9dba"）
            video_path: 对应的视频文件路径（如 "data/video/movie.mp4"）
        """
        logging.info(f"开始重新提取音频...")
        logging.info(f"处理IP目录: {ip_dir}")
        logging.info(f"使用视频文件: {video_path}")

        # 验证视频文件是否存在
        if not Path(video_path).exists():
            logging.error(f"视频文件不存在: {video_path}")
            return False

        # 构建数据集文件路径
        dataset_file = Path("output") / ip_dir / f"word_dataset_{ip_dir}_with_style.json"

        if not dataset_file.exists():
            logging.error(f"数据集文件不存在: {dataset_file}")
            return False

        logging.info(f"找到数据集文件: {dataset_file}")

        # 处理数据集文件
        success = self.process_dataset_file(dataset_file, video_path)

        if success:
            logging.info("音频重新提取完成！")
        else:
            logging.error("音频重新提取失败")

        return success

def main():
    """主函数"""
    # ========== 在这里配置要处理的IP目录和视频文件 ==========

    # 示例配置 - 请根据实际情况修改
    ip_dir = "IP1fef9dba"  # 要处理的IP目录名
    video_path = "data/video/zootopia.mp4"  # 对应的视频文件路径（支持.mp4, .mkv等格式）

    # 更多示例：
    # ip_dir = "IP1750926257"
    # video_path = "data/video/kungfu_panda.mkv"

    # ip_dir = "IP9754cf1e"
    # video_path = "data/video/peppa_pig.mp4"

    # =====================================================

    logging.info(f"配置信息:")
    logging.info(f"  IP目录: {ip_dir}")
    logging.info(f"  视频文件: {video_path}")

    extractor = AudioReExtractor()
    success = extractor.run(ip_dir, video_path)

    if success:
        logging.info("音频重新提取完成！")
        sys.exit(0)
    else:
        logging.error("音频重新提取过程中出现错误")
        sys.exit(1)

if __name__ == "__main__":
    main()
